import matplotlib.pyplot as plt
import torch
from stylegan2_pytorch import Generator

# 确保你的设备是可用的GPU，如果不是，则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 加载预训练的StyleGAN模型
# 请替换为你的预训练模型路径
model_path = 'path_to_your_pretrained_stylegan_model.pth'
# 假设是StyleGAN2，生成256x256的图片，需要根据你的模型调整参数
g_ema = Generator(256, 512, 8).to(device)
g_ema.load_state_dict(torch.load(model_path, map_location=device)['g_ema'])

# 生成随机噪声
random_seed = torch.randn(1, 512).to(device)  # 512是StyleGAN2的latent dimension

# 生成图片
with torch.no_grad():
    g_ema.eval()
    fake_img = g_ema([random_seed])

# 将生成的图片转换为可显示的格式
fake_img = (fake_img + 1) / 2.0  # 反标准化到[0, 1]
fake_img = fake_img.clamp(0, 1)  # 限制值在[0, 1]范围内
fake_img = fake_img.cpu().numpy().transpose(0, 2, 3, 1)  # 转换为HWC格式

# 保存生成的图片
plt.imsave('generated_image.png', fake_img[0])
